{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import copy\n",
    "import numpy as np\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def draw_fill(puzzle, patternLength, patternWidth, start, count, solList):\n",
    "    count += 1\n",
    "    puzzleLength, puzzleWidth = puzzle.shape\n",
    "    patternNum = (puzzleWidth*puzzleLength)/(patternWidth*patternLength)\n",
    "    \n",
    "    horizonal = False\n",
    "    if start[0] + patternLength <= puzzleLength and start[1] + patternWidth <= puzzleWidth:\n",
    "        horizonal = True\n",
    "        #if (puzzle[start[0]:start[0]+patternLength, start[1]:start[1]+patternWidth] != 0).any():\n",
    "        for i in range(start[0], start[0]+patternLength):\n",
    "             for j in range(start[1], start[1]+patternWidth):\n",
    "                 if puzzle[i][j] != 0:\n",
    "                     horizonal = False\n",
    "    if horizonal:\n",
    "        newPuzzle = copy.deepcopy(puzzle)\n",
    "        for i in range(start[0], start[0]+patternLength):\n",
    "            for j in range(start[1], start[1]+patternWidth):\n",
    "                newPuzzle[i][j] = count\n",
    "        if count == patternNum:\n",
    "            solList.append(newPuzzle)\n",
    "            return\n",
    "        for i in range(start[0], puzzleLength):\n",
    "            for j in range(0, puzzleWidth):\n",
    "                if newPuzzle[i][j] == 0:\n",
    "                    newStart = (i, j)\n",
    "                    break\n",
    "            else:\n",
    "                continue\n",
    "            break\n",
    "        draw_fill(newPuzzle, patternLength, patternWidth, newStart, count, solList)\n",
    "\n",
    "    vertical = False\n",
    "    if patternLength != patternWidth and start[0]+patternWidth <= puzzleLength and start[1]+patternLength <= puzzleWidth:\n",
    "        vertical = True\n",
    "        for i in range(start[0], start[0]+patternWidth):\n",
    "            for j in range(start[1], start[1]+patternLength):\n",
    "                if puzzle[i][j] != 0:\n",
    "                    vertical = False\n",
    "    if vertical:\n",
    "        newPuzzle = copy.deepcopy(puzzle)\n",
    "        for i in range(start[0], start[0]+patternWidth):\n",
    "            for j in range(start[1], start[1]+patternLength):\n",
    "                newPuzzle[i][j] = count\n",
    "        if count == patternNum:\n",
    "            solList.append(newPuzzle)\n",
    "            return\n",
    "        for i in range(start[0], puzzleLength):\n",
    "            for j in range(0, puzzleWidth):\n",
    "                if newPuzzle[i][j] == 0:\n",
    "                    newStart = (i, j)\n",
    "                    break\n",
    "            else:\n",
    "                continue\n",
    "            break\n",
    "        draw_fill(newPuzzle, patternLength, patternWidth, newStart, count, solList)\n",
    "\n",
    "def backtrack(puzzleLength, puzzleWidth, patternLength, patternWidth):\n",
    "    patternNum = (puzzleWidth*puzzleLength)/(patternWidth*patternLength)\n",
    "    solList = []\n",
    "    if patternNum%1 == 0:\n",
    "        inputPuzzle = np.zeros((puzzleLength, puzzleWidth))\n",
    "        draw_fill(inputPuzzle, patternLength, patternWidth, (0, 0), 0, solList)\n",
    "    #solList = np.asarray(solList).reshape((puzzleLength, puzzleWidth))\n",
    "    return solList"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([array([[1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.],\n",
       "         [1., 1., 1., 1.]]),\n",
       "  array([[1., 1., 1., 1.],\n",
       "         [2., 2., 2., 2.],\n",
       "         [3., 3., 3., 3.],\n",
       "         [4., 4., 4., 4.],\n",
       "         [5., 5., 5., 5.],\n",
       "         [6., 6., 6., 6.],\n",
       "         [7., 7., 7., 7.],\n",
       "         [8., 8., 8., 8.]])],\n",
       " array([1., 1.]))"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_cost_c(conf, L, cluster_info=None):\n",
    "    # homogeneous setting; in real setting, we access cluster to get cost_c\n",
    "    num_stages = int(np.max(conf))\n",
    "    stage_cost = []\n",
    "    for i in range(1, num_stages):\n",
    "        b = np.where(conf == i)\n",
    "        c = np.where(conf == i+1)\n",
    "        # All pairs of GPU in the same node\n",
    "        if (b[1] == c[1]).all():\n",
    "            stage_cost.append(0)\n",
    "        else:\n",
    "            stage_cost.append(1)\n",
    "    stage_cost = np.asarray(stage_cost).reshape((1,-1))\n",
    "    ret = copy.deepcopy(stage_cost)\n",
    "    for i in range(L-1):\n",
    "        ret = np.concatenate((ret, stage_cost), axis=0)\n",
    "    return ret\n",
    "\n",
    "def get_cost_e(conf, L, cluster_info=None):\n",
    "    # homogeneous setting; in real setting, we access cluster to get cost_e\n",
    "    # return amp_simulator()\n",
    "    #print(conf.shape[0] * conf.shape[1])\n",
    "    num_gpus_per_pipeline = conf.shape[0] * conf.shape[1] / np.max(conf)\n",
    "    return np.ones(L) / num_gpus_per_pipeline\n",
    "\n",
    "def generate_initial(M, N, threads=2):\n",
    "    h_w_list = []\n",
    "    \n",
    "    h_w_list.append((M, 1))\n",
    "    h_w_list.append((1, N))\n",
    "    known = {}\n",
    "    \n",
    "    configs = []\n",
    "    for (h, w) in h_w_list:\n",
    "        solution = backtrack(M, N, h, w)\n",
    "        \n",
    "        assert len(solution) > 0\n",
    "        config_idx = np.random.choice(len(solution), size=1)[0]\n",
    "        config = solution[config_idx]\n",
    "        configs.append(config)\n",
    "        \n",
    "        solution.pop(config_idx)\n",
    "        \n",
    "        known[(h, w)] = solution\n",
    "        \n",
    "    #print(np.asarray(configs[0]))\n",
    "    return h_w_list, configs, known\n",
    "    \n",
    "\n",
    "def cool_down(iter, max_iter, init_temp):\n",
    "    return init_temp * (1 - iter / max_iter)\n",
    "\n",
    "def neighbor(cur, known, M, N, maximum_try = 10):\n",
    "    h, w = cur\n",
    "    \n",
    "    time_s = time.time()\n",
    "    while time.time() - time_s < 10:\n",
    "        index = np.random.choice([0,1], size=1)[0]\n",
    "        if index == 0:\n",
    "            valid = []\n",
    "            upper = min(M, N)\n",
    "            upper = min((M*N) // w, upper) + 1\n",
    "            \n",
    "            for i in range(1, upper):\n",
    "                if (i, w) in known.keys():\n",
    "                    solution = known[(i, w)]\n",
    "                else:\n",
    "                    solution = backtrack(M, N, i, w)\n",
    "                    known[(i, w)] = solution\n",
    "\n",
    "                if len(solution) > 0:\n",
    "                    valid.append(i)\n",
    "\n",
    "            if len(valid) == 0:\n",
    "                continue\n",
    "                #return\n",
    "                \n",
    "            new_h = np.random.choice(valid, size=1)[0]\n",
    "            \n",
    "            # TODO\n",
    "            new_config_idx = np.random.choice(len(known[(new_h, w)]), size=1)[0]\n",
    "            ret = known[(new_h, w)].pop(new_config_idx)\n",
    "            return new_h, w, ret\n",
    "\n",
    "        else:\n",
    "            valid = []\n",
    "            upper = min(M, N)\n",
    "            upper = min((M*N) // h, upper) + 1\n",
    "            for i in range(1, upper):\n",
    "                if (h, i) in known.keys():\n",
    "                    solution = known[(h, i)]\n",
    "                else:\n",
    "                    solution = backtrack(M, N, h, i)\n",
    "                    known[(h, i)] = solution\n",
    "\n",
    "                if len(solution) > 0:\n",
    "                    valid.append(i)\n",
    "\n",
    "            if len(valid) == 0:\n",
    "                continue\n",
    "\n",
    "            new_w = np.random.choice(valid, size=1)[0]\n",
    "            new_config_idx = np.random.choice(len(known[(h, new_w)]), size=1)[0]\n",
    "            ret = known[(h, new_w)].pop(new_config_idx)  \n",
    "            return h, new_w, ret\n",
    "    return None\n",
    "    \n",
    "def predict(configs, L, B):\n",
    "    costs = []\n",
    "    for i in range(len(configs)):\n",
    "        config = configs[i]\n",
    "        config = np.asarray(config)\n",
    "        #config = config.reshape((config.shape[0], config.shape[2]))\n",
    "        cost_e = get_cost_e(config, L)\n",
    "        cost_c = get_cost_c(config, L)\n",
    "        k = int(np.max(config))\n",
    "\n",
    "        # refer to pipeling slicing\n",
    "        cost = pipe_dp(L, cost_e, cost_c, k, B)[1]\n",
    "        costs.append(cost)\n",
    "    return np.asarray(costs)\n",
    "\n",
    "# number of GPU per node\n",
    "M = 8\n",
    "# \n",
    "N = 4\n",
    "num_iter = 500\n",
    "init_t = 1\n",
    "\n",
    "# 16 layers network, 3 microbatches\n",
    "L = 16\n",
    "B = 3\n",
    "\n",
    "h_w, configs, known = generate_initial(M, N)\n",
    "costs = predict(configs, L, B)\n",
    "\n",
    "for i in range(num_iter):\n",
    "    cur_t = cool_down(i, num_iter, init_t)   \n",
    "    \n",
    "    new_configs = []\n",
    "    new_h_w = []\n",
    "    \n",
    "    for (h, w) in h_w:\n",
    "        step = neighbor((h, w), known, M, N)\n",
    "        if step is None:\n",
    "            new_h, new_w, new_config = (h, w, configs[h_w.index((h,w))])\n",
    "            \n",
    "        else:\n",
    "            new_h, new_w, new_config = step\n",
    "        if step is None:\n",
    "            assert False\n",
    "        else:\n",
    "            pass\n",
    "            #print(step)\n",
    "        new_h_w.append((new_h, new_w))\n",
    "        new_configs.append(new_config)\n",
    "        \n",
    "    new_costs = predict(new_configs, L, B)\n",
    "    \n",
    "    acc_prob = np.exp(np.minimum((costs - new_costs)/ (cur_t+1e-5) , 0))\n",
    "    \n",
    "    acc_index = (np.random.random(len(acc_prob)) < acc_prob)\n",
    "    \n",
    "    for j in range(len(configs)):\n",
    "        if acc_index[j]:\n",
    "            configs[j] = new_configs[j]\n",
    "            costs[j] = new_costs[j]\n",
    "\n",
    "configs, costs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scratch code below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def placement_reachable(M, N, m, n, s_joint):\n",
    "    #horizontal_tile = np.asarray(list(range(m * n))).reshape((m, n))\n",
    "    #vertical_tile = np.transpose(horizontal_tile)\n",
    "    horizontal_tile = np.ones((m,n))\n",
    "    vertical_tile = np.ones((n,m))\n",
    "    vertical_tile[0] = 0\n",
    "    \n",
    "    t = True\n",
    "    i = 0\n",
    "    while i < N:\n",
    "        match = False\n",
    "        # Check whether horizontal \n",
    "        if i <= N - n:\n",
    "            for j in range(n-m, n):\n",
    "                #print(s_joint[j:, i:i+n])\n",
    "                match_height = n-j\n",
    "   #            print(match_height)\n",
    "                if (s_joint[j:, i:i+n] == horizontal_tile[:match_height,:]).all():\n",
    "   #                 print(i, j, \"h\", s_joint[j:, i:i+n], horizontal_tile[:match_height,:], match_height)\n",
    "                    i += n\n",
    "                    if j != n-m:\n",
    "                        t = False\n",
    "                    match = True\n",
    "                    break\n",
    "        \n",
    "        if i <= N - m:\n",
    "            for j in range(n):\n",
    "                #print(s_joint,j,i,m, s_joint[j:, i:i+m])\n",
    "                match_height = n-j\n",
    "                if (s_joint[j:, i:i+m] == vertical_tile[:match_height,:]).all():\n",
    "  #                  print(i, j, \"v\", s_joint[j:, i:i+n], vertical_tile[:match_height,:], match_height)\n",
    "                    i += m\n",
    "                    if j != 0:\n",
    "                        t = False\n",
    "                    match = True\n",
    "                    break\n",
    "        \n",
    "        if not match:\n",
    "            return False, _\n",
    "    return True, t\n",
    "\n",
    "# ! Always assume m < n\n",
    "def init(M, N, m, n, s_array):\n",
    "    h, w = s_array.shape\n",
    "    checked = np.zeros((h, w))\n",
    "    i = 0\n",
    "    j = 0\n",
    "#     horizontal_tile = np.asarray(list(range(m * n))).reshape((m, n))\n",
    "#     vertical_tile = np.transpose(horizontal_tile)\n",
    "    horizontal_tile = np.ones((m,n))\n",
    "    vertical_tile = np.ones((n,m))\n",
    "    vertical_tile[0] = 0\n",
    "    \n",
    "    \n",
    "    #print(s_array)\n",
    "    terminate = True\n",
    "    for i in range(h):\n",
    "        for j in range(w):\n",
    "            if checked[i][j] == 1:\n",
    "                continue\n",
    "                \n",
    "            # Check horizontal\n",
    "            if i <= M - m and j <= N - n:\n",
    "                match_height = min(h-i, m)\n",
    "                if (s_array[i:i+match_height, j:j+n] == horizontal_tile[:match_height,:]).all() and (checked[i:i+m, j:j+n] != 1).all():\n",
    "                    checked[i:i + m, j: j + n] = 1\n",
    "                    if match_height != m:\n",
    "                        terminate = False\n",
    "                    continue\n",
    "            \n",
    "            # Check vertical\n",
    "            if i <= M - n and j <= N - m:\n",
    "                match_height = min(h-i, n)\n",
    "                if (s_array[i:i+match_height, j:j+m] == vertical_tile[:match_height,:]).all() and (checked[i:i+n, j:j+m] != 1).all():\n",
    "                    checked[i:i + n, j: j + m] = 1\n",
    "                    if match_height != n:\n",
    "                        terminate = False\n",
    "                    continue\n",
    "            #print(i, j, s_array, checked)\n",
    "            return False, _\n",
    "    return True, terminate\n",
    "        \n",
    "# returns all possible pipe group configurations\n",
    "def generate_placement(grid, len_1, len_2):\n",
    "    tot_len = len_1 * len_2\n",
    "    # possible configuration number for a row\n",
    "    from itertools import product\n",
    "    #possible_s = list(product(range(tot_len),repeat = grid.shape[1]*(len_2-1)))\n",
    "    #single_possible_s = list(product(list(range(tot_len)),repeat = grid.shape[1]))\n",
    "    \n",
    "    possible_s = list(product(range(2),repeat = grid.shape[1]*(len_2-1)))\n",
    "    single_possible_s = list(product(list(range(2)),repeat = grid.shape[1]))\n",
    "    \n",
    "    #print(possible_s, single_possible_s)\n",
    "    for i in range(len(possible_s)):\n",
    "        possible_s[i] = np.asarray(list(possible_s[i])).reshape(1,-1)\n",
    "    \n",
    "    for i in range(len(single_possible_s)):\n",
    "        single_possible_s[i] = np.asarray(list(single_possible_s[i])).reshape(1,-1)\n",
    "    \n",
    "    \n",
    "    # the solution will be the union of all possible configurations in the last row\n",
    "    dp = [[None for j in range(len(possible_s))] for i in range(grid.shape[0])]\n",
    "    \n",
    "    # initialize the first (len_1 -1) row\n",
    "    for s_index in range(len(possible_s)):\n",
    "        valid, terminate = init(grid.shape[0], grid.shape[1], len_1, len_2, possible_s[s_index].reshape(-1, grid.shape[1]))\n",
    "        if valid:\n",
    "            dp[0][s_index] =  [(possible_s[s_index].reshape(-1, grid.shape[1]), terminate)]\n",
    "            #print(possible_s[s_index])\n",
    "    print(dp[0])\n",
    "    # dp by row index\n",
    "    for i in range(len_2-1, grid.shape[0]):\n",
    "        print(\" \")\n",
    "        print(dp[i-1], i)\n",
    "        print(\" \")\n",
    "        # iterate through all possibly reachable row?\n",
    "        #j = i - 1\n",
    "        for s_index_1 in range(len(possible_s)):\n",
    "      #      print(\"haha\", s_index_1, len(possible_s))\n",
    "            for s_index_2 in range(len(single_possible_s)):\n",
    "                s_1 = possible_s[s_index_1]\n",
    "                s_2 = single_possible_s[s_index_2]\n",
    "               # print(s_1, s_2)\n",
    "                s_joint = np.concatenate((s_1, s_2), axis=0)\n",
    "                # early return if the last rows themselves are not possible\n",
    "                #print(s_joint, valid)\n",
    "                if dp[i-1][s_index_1] is None:\n",
    "                    print(i-1, s_index_1)\n",
    "                    continue\n",
    "                    \n",
    "                #valid, terminate =  placement_reachable(grid.shape[0], grid.shape[1], len_1, len_2, s_joint)\n",
    "                #valid, terminate =  init(grid.shape[0], grid.shape[1], len_1, len_2, s_joint)\n",
    "                valid, terminate =  placement_reachable(grid.shape[0], grid.shape[1], len_1, len_2, s_joint)\n",
    "     #           print(s_joint, valid)\n",
    "                if valid:\n",
    "                    if dp[i][s_index_2] is None:\n",
    "                        dp[i][s_index_2] = []\n",
    "                    for solution in dp[i-1][s_index_1]:\n",
    "                        #print(i-1,solution)\n",
    "                        sol, _ = solution\n",
    "                        s_joint_sol = np.concatenate((copy.deepcopy(sol), s_2), axis=0)\n",
    "                        dp[i][s_index_2].append((s_joint_sol, terminate))\n",
    "#     print(dp[0])\n",
    "#     print(dp[1])\n",
    "#     print(dp[2])\n",
    "    ret_sol = []\n",
    "    for i in range(len(single_possible_s)):\n",
    "        s = possible_s[i]\n",
    "        if dp[grid.shape[0]-1][i] is None:\n",
    "            continue\n",
    "        for (sol, t) in dp[grid.shape[0]-1][i]:\n",
    "            if t:\n",
    "                ret_sol.append(sol)\n",
    "    return ret_sol\n",
    "\n",
    "# for len_1 in factors:\n",
    "#     # Genarate all possible configuratinos\n",
    "#     remain = num_gpu / len_1\n",
    "#     factors_2 = []\n",
    "#     for i in range(1, min(cluster_shape) + 1):\n",
    "#         if remain % i == 0:\n",
    "#             factors_2.append(i)\n",
    "#         for len_2 in factors_2:\n",
    "#             num_cut = num_gpu / (len_1*len_2)\n",
    "#             confs = generate_placement(grid, len_1. len_2)\n",
    "#             for conf in confs:\n",
    "#                 cost_c = get_cost_c(conf)\n",
    "#                 cost_e = get_cost_e(conf)\n",
    "#                 opt_pipe = pipe_dp(L, cost_e, cost_c, num_cut, B)\n",
    "#                 cost = amp_simulator(conf, opt_pipe)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}